from flask import Flask, request, jsonify
import os
import json
from typing import Dict, List, Any, Optional
import argparse


class EnhancedMCPServer:
    """
    An enhanced Model Context Protocol server that supports dynamic tool selection
    and provides additional context based on user queries.
    """

    def __init__(self, data_dir: str = "./data"):
        """
        Initialize the MCP server.
        """
        self.data_dir = data_dir
        self.knowledge_base = self._load_knowledge_base()
        self.tools = self._register_tools()

    def _load_knowledge_base(self) -> Dict[str, Any]:
        """
        Load the knowledge base from files in the data directory.
        """
        knowledge_base = {}

        # Create the data directory if it doesn't exist
        os.makedirs(self.data_dir, exist_ok=True)

        # Check for knowledge base file
        kb_path = os.path.join(self.data_dir, "knowledge_base.json")

        if os.path.exists(kb_path):
            try:
                with open(kb_path, 'r') as f:
                    knowledge_base = json.load(f)
            except json.JSONDecodeError:
                print(f"Error loading knowledge base from {kb_path}")
                knowledge_base = {}
        else:
            # Create a sample knowledge base if none exists
            sample_knowledge = {
                "model_context_protocol": {
                    "title": "Model Context Protocol (MCP)",
                    "description": "A protocol for providing additional context to language models.",
                    "details": [
                        "MCP enables models to access external data sources.",
                        "It provides a standardized way to augment model inputs with relevant context.",
                        "MCP servers can perform retrieval and processing before the model generates a response.",
                        "The protocol is designed to be model-agnostic and work with any LLM system.",
                        "It supports tools that can be dynamically selected by the LLM."
                    ],
                    "use_cases": [
                        "Document retrieval",
                        "Knowledge base access",
                        "Database querying",
                        "API integration",
                        "Tool selection and usage"
                    ]
                },
                "python_programming": {
                    "title": "Python Programming Language",
                    "description": "Python is a high-level, interpreted programming language known for its readability and versatility.",
                    "key_features": [
                        "Simple, easy-to-learn syntax",
                        "Interpreted language",
                        "Dynamically typed",
                        "Object-oriented",
                        "Extensive standard library",
                        "Wide range of third-party packages"
                    ],
                    "common_uses": [
                        "Web development",
                        "Data science and analysis",
                        "Machine learning and AI",
                        "Automation and scripting",
                        "Scientific computing"
                    ]
                }
            }

            # Save the sample knowledge base
            with open(kb_path, 'w') as f:
                json.dump(sample_knowledge, f, indent=2)

            knowledge_base = sample_knowledge
            print(f"Created sample knowledge base at {kb_path}")

        return knowledge_base

    def _register_tools(self) -> List[Dict[str, Any]]:
        """
        Register available tools with descriptions.

        In a production system, this might be more dynamic or pluggable.
        """
        return [
            {
                "name": "knowledge_retrieval",
                "description": "Retrieve information from the knowledge base on a specific topic",
                "parameters": ["topic"]
            },
            {
                "name": "topic_search",
                "description": "Search across the knowledge base for information related to keywords",
                "parameters": ["keywords"]
            }
        ]

    def get_tools(self) -> List[Dict[str, Any]]:
        """
        Return the list of available tools.
        """
        return self.tools

    def process_query(
            self,
            query: str,
            tool: Optional[str] = None,
            tool_params: Optional[Dict[str, Any]] = None,
            history: Optional[List[Dict[str, str]]] = None
    ) -> Dict[str, Any]:
        """
        Process a query and return relevant context.
        Optionally use a specific tool if specified.
        """
        if history is None:
            history = []

        if tool_params is None:
            tool_params = {}

        # If a specific tool is requested, use it
        if tool and tool in [t["name"] for t in self.tools]:
            if tool == "knowledge_retrieval":
                return self._knowledge_retrieval_tool(query, tool_params)
            elif tool == "topic_search":
                return self._topic_search_tool(query, tool_params)

        # Otherwise, fall back to simple keyword matching
        return self._default_retrieval(query)

    def _knowledge_retrieval_tool(self, query: str, params: Dict[str, Any]) -> Dict[str, Any]:
        """
        Tool for retrieving specific topics from the knowledge base.
        """
        topic = params.get("topic", "").lower()

        # Look for the topic in knowledge base keys
        for key in self.knowledge_base:
            if topic in key.lower():
                kb_entry = self.knowledge_base[key]

                # Format the context
                context = f"# {kb_entry.get('title', key.capitalize())}\n\n"
                context += f"{kb_entry.get('description', '')}\n\n"

                # Add any additional fields from the knowledge base entry
                for field, value in kb_entry.items():
                    if field not in ['title', 'description'] and isinstance(value, list):
                        context += f"## {field.replace('_', ' ').capitalize()}\n\n"
                        for item in value:
                            context += f"- {item}\n"
                        context += "\n"

                return {
                    "context": context,
                    "context_used": True,
                    "tool_used": "knowledge_retrieval",
                    "topic_found": key
                }

        # If no match found
        return {
            "context": f"No information found for topic: {topic}",
            "context_used": False,
            "tool_used": "knowledge_retrieval",
            "topic_found": None
        }

    def _topic_search_tool(self, query: str, params: Dict[str, Any]) -> Dict[str, Any]:
        """
        Tool for searching across the knowledge base using keywords.
        """
        keywords = params.get("keywords", [])
        if isinstance(keywords, str):
            keywords = [keywords]

        if not keywords:
            # Extract potential keywords from the query itself
            # This is a simple implementation - in production, you might use NLP
            query_words = query.lower().split()
            keywords = [word for word in query_words if len(word) > 3]

        results = []

        # Search through all knowledge base entries
        for key, entry in self.knowledge_base.items():
            # Check title and description
            entry_text = f"{entry.get('title', '')} {entry.get('description', '')}".lower()

            # Include all list fields
            for field, value in entry.items():
                if isinstance(value, list):
                    entry_text += " " + " ".join(value).lower()

            # Check if any keyword matches
            if any(keyword.lower() in entry_text for keyword in keywords):
                results.append({
                    "topic": key,
                    "title": entry.get("title", key.capitalize()),
                    "description": entry.get("description", "")
                })

        if results:
            # Format the context with search results
            context = "# Search Results\n\n"

            for result in results:
                context += f"## {result['title']}\n"
                context += f"{result['description']}\n\n"

            return {
                "context": context,
                "context_used": True,
                "tool_used": "topic_search",
                "matches": len(results),
                "topics_found": [r["topic"] for r in results]
            }

        return {
            "context": f"No information found for keywords: {', '.join(keywords)}",
            "context_used": False,
            "tool_used": "topic_search",
            "matches": 0,
            "topics_found": []
        }

    def _default_retrieval(self, query: str) -> Dict[str, Any]:
        """
        Default retrieval method using simple keyword matching.
        """
        context = ""
        context_used = False
        found_topics = []

        # Very simple keyword-based retrieval
        for key, entry in self.knowledge_base.items():
            # Check if any keywords from the topic appear in the query
            if key.lower().replace("_", " ") in query.lower():
                # Format the entry
                if context:
                    context += "\n\n----------\n\n"  # Separator between topics

                context += f"# {entry.get('title', key.capitalize())}\n\n"
                context += f"{entry.get('description', '')}\n\n"

                # Add any additional fields from the knowledge base entry
                for field, value in entry.items():
                    if field not in ['title', 'description'] and isinstance(value, list):
                        context += f"## {field.replace('_', ' ').capitalize()}\n\n"
                        for item in value:
                            context += f"- {item}\n"
                        context += "\n"

                context_used = True
                found_topics.append(key)

        return {
            "context": context,
            "context_used": context_used,
            "tool_used": "default_retrieval",
            "topics_found": found_topics
        }

    def add_to_knowledge_base(self, key: str, data: Dict[str, Any]) -> None:
        """
        Add new information to the knowledge base.
        """
        self.knowledge_base[key] = data

        # Save the updated knowledge base
        kb_path = os.path.join(self.data_dir, "knowledge_base.json")
        with open(kb_path, 'w') as f:
            json.dump(self.knowledge_base, f, indent=2)


# Flask application
app = Flask(__name__)
mcp_server = None


@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint."""
    return jsonify({"status": "healthy"})


@app.route('/tools', methods=['GET'])
def get_tools():
    """Get available tools endpoint."""
    return jsonify({"tools": mcp_server.get_tools()})


@app.route('/query', methods=['POST'])
def process_query():
    """Process a query and return relevant context."""
    try:
        data = request.json

        if not data or 'query' not in data:
            return jsonify({"error": "Missing query parameter"}), 400

        query = data['query']
        history = data.get('history', [])
        tool = data.get('tool')
        tool_params = data.get('tool_params', {})

        result = mcp_server.process_query(query, tool, tool_params, history)

        return jsonify(result)

    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/add', methods=['POST'])
def add_knowledge():
    """Add new knowledge to the knowledge base."""
    try:
        data = request.json

        if not data or 'key' not in data or 'data' not in data:
            return jsonify({"error": "Missing key or data parameters"}), 400

        key = data['key']
        knowledge_data = data['data']

        mcp_server.add_to_knowledge_base(key, knowledge_data)

        return jsonify({"status": "success", "message": f"Added knowledge for key: {key}"})

    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.after_request
def after(response):
    # todo with response
    print(response.status)
    print(response.headers)
    print(response.get_data())
    return response


def main():
    """Run the MCP server."""
    parser = argparse.ArgumentParser(description="Enhanced Model Context Protocol Server")
    parser.add_argument("--host", default="localhost", help="Host to run the server on")
    parser.add_argument("--port", default=3000, type=int, help="Port to run the server on")
    parser.add_argument("--data-dir", default="./data", help="Directory for knowledge base data")

    args = parser.parse_args()

    global mcp_server
    mcp_server = EnhancedMCPServer(data_dir=args.data_dir)

    print(f"Starting Enhanced MCP server on {args.host}:{args.port}")
    app.run(host=args.host, port=args.port, debug=True)


if __name__ == "__main__":
    main()
